open Glob
open Histo
open Plplot
open Printf
(** Draw histograms using Plplot *)
(** Основной стиль *)
let def_style ?(color=1) () =
    plcol0 color;
    plwid 2;
    plsmin 0.2 1.;
    ()

let vp_y_min=0.19
let vp_y_max=0.91
let vp_x_min=0.12
let vp_x_max=0.95

let next_page x1 x2 y1 y2=
    plclear ();
    plvpor vp_x_min vp_x_max vp_y_min vp_y_max;
    plwind x1 x2 y1 y2;
    ()

let make_logarithmic x min_x =
    let nx=Array.length x in
    Array.init nx (fun i ->
        if x.(i) > min_x
        then log10(x.(i))
        else log10(min_x) ) 

let cmap1_init () =
  let i = [|0.0; 1.0|] in
  let h = [|0.0; 240.0|] in (* blue -> green -> yellow -> red *)
  let l = [|0.6; 0.6|] in
  let s = [|0.8; 0.8|] in
  plscmap1n 256;
  plscmap1l false i h l s None;
  ()
;;

type place = Line1 | Line2

let set_2d_palette = function
    2 -> cmap1_init () 
  | 1 -> plspal1 "cmap1_gray.pal" true
  | _ -> failwith "Unkown palette"

 let format_float v=
     let s1=Printf.sprintf "%.3g" v in
     let s2=Str.global_replace (Str.regexp "e\+0\(.\)") "⋅10#u\1#d" s1 in
     let s3=Str.global_replace (Str.regexp "e\+\(..\)") "⋅10#u\1#d" s2 in
     s3

 (** draws line with statistics information about histogram;
     function list_info returns list of pairs (name, value), 
     e.g. ("Entries", 50.0) *)
 let draw_stats_for ?(place=Line1) ~list_info hists=
     let list_list_info=List.map list_info hists in
     let h1=List.hd list_list_info in
     let start_list=List.map (fun (name,value) -> (name,[])) h1 in
     let all_info = List.fold_left (fun all_info info1 ->
         List.map2 (fun (name,list) (name,value) -> 
             (name,value::list)
         ) all_info info1
     ) start_list list_list_info 
     in
     let info_string = List.fold_left (fun str (name,value_list) ->
         let l = List.map format_float value_list in
         str^name^(List.hd l)^(
             List.fold_left (fun s n -> s ^ "," ^ n) "" (List.tl l)
         )^"; "
     ) "" all_info 
     in
     plvpor 0.0 1. 0. 0.01;
     let y = match place with Line1 -> 0.01 | Line2 -> 0.1 in
     plmtex "t" 1. 0.05 y info_string
;;
let rec list_iter2_opt ~last_opt f l lopt =
    ( match l with
      | [] -> ()
      | a::l -> 
           ( match lopt with
             | [] -> 
                     f a last_opt; list_iter2_opt ~last_opt f l []
             | a2::lopt -> 
                     f a a2; list_iter2_opt ~last_opt:a2 f l lopt
           )
    ) 
;;
module H1_pl (*: H1_draw*) = struct
 open Histo.H1
 let draw_err (x,y,y_err_min,y_err_max) =
     plerry x y_err_min y_err_max

 let draw_bin (x,y) =
     plbin x y []

 let draw_line (x,y) last_x=
     let l=Array.length y in
     let x_dots=Array.init (2*l) ( fun i ->
         let k=(i+1)/2 in
         if k<l
         then x.(k)
         else last_x
     ) in
     let y_dots=Array.init (2*l) ( fun i ->
         y.( i/2 )
     ) in
     plline x_dots y_dots;
     ()

 let draw_polyline (x,y) =
     plline x y;
     ()

 let list_info h=
     let (_,mean,rms)=max_mean_rms h in
     [
         ("N=", Int64.to_float h.entries);
         ("∑#dx∈r#uy=", integral h);
         (*("∑#dx∉r#uy=", sum_out h);*)
         ("∑#d<#uy=", h.underflow);
         ("∑#d>#uy=", h.overflow);
         ("Mx=", mean);
         ("√Dx=", rms);
     ]

 let is_opt_log log_opt = function 
     X -> List.mem Log_X log_opt 
  |  Y -> List.mem Log_Y log_opt

 let box_parameters plot_opt proj =
     let is_log axis = is_opt_log plot_opt.logscale axis in
     (match proj with X -> "bcnst" | Y -> "bcnstv")^
     (if is_log proj then "l" else "")^
     (if is_grid plot_opt proj then "gh" else "")

 type dot_pos = Left | Center

(* let draw_hist ?(draw_opt=standard_plot) ?(add_hists=[]) 
        ?(add_hists_opt=[]) ?(add_hists_fit=[]) h=*)
  let draw 
        ?(plot_opt=standard_plot) ?(opt=[]) ?(sec_hists=[]) ?(sec_opt=[])
        hists =
     def_style();
     (match hists with 
      | [] -> ()
      | h :: other_hists ->  (* 1st histogram is main one *)
         let draw_opt_1, draw_opt_other = ( match opt with 
             [] -> (standard_show, []) | opt1::ol -> (opt1, ol)) in
         let is_log axis = is_opt_log plot_opt.logscale axis in
         let xopt = box_parameters plot_opt X and
         yopt = box_parameters plot_opt Y in
(*         print_endline ("xopt :"^xopt);*)
(*         let xopt="bcnst"  ^(if is_log X then "l" else "")
         and yopt="bcnstv" ^(if is_log Y then "l" else "") in*)
         let x_dots=H1.x_dots h in 
         let get_minmax axis arr =
             if is_log axis
             then (
                 try
                     let min, max = Glob.min0maxf arr in
                     let add=0.5*.(max-.min) in
                     (0.9*.min),(max+.add)
                 with Not_found ->
                     raise (Invalid_histogram ("there are no positive"^
                     " or finite values in histogram "^h.H1.info.name))
             ) else (
                 if axis = X
                 then Axis.bounds h.xa
                 else (
                     let min,max = Glob.minfmaxf arr in
                     let add=0.05*.(max-.min) in
                     let min=
                         if min>= 0. && min-.add<=0. then 0. else (min-.add)
                     in 
                     min,(max+.add)
                 )
             )
         in
         let get axis arr axis_min=
             if is_log axis
             then (make_logarithmic arr axis_min)
             else arr
         in
         let x_min,x_max=get_minmax X x_dots in
    (*     let x_max=last_x in*)
         let h1_y_min,h1_y_max=get_minmax Y h.H1.arr in
         let y_min,y_max=List.fold_left (fun (y1min,y1max) h2 ->
             let (y2min,y2max) = get_minmax Y h2.H1.arr in
             (min y1min y2min, max y1max y2max)
         ) (h1_y_min,h1_y_max) other_hists in
         let try_log axis a b = 
             if is_log axis
             then (log10 a),(log10 b)
             else a,b
         in
         let x1,x2=try_log X x_min x_max
         and y1,y2=try_log Y y_min y_max in
         next_page x1 x2 y1 y2;
         let wx,wy,n = plcalc_world 0.5 0.5 in 
         printf "!!!%f %f %i\n" wx wy n; flush stdout;
         plbox xopt 0.0 0 yopt 0.0 0;
         pllab (Axis.name h.xa) h.y_title h.info.title;
         let get_dots pos h =
             let dots = match pos with
                  Left -> H1.x_left_dots h
                | Center -> H1.x_dots h in
             let x=get X dots x_min
             and y=get Y h.H1.arr y_min in
             let err i yi = match h.sq_errors with
                    Some a -> a.(i)
                  | None -> sqrt(yi)
             in
             let y_err_min = get Y 
                     (Array.mapi (fun i yi -> yi -. (err i yi)) h.H1.arr) y_min in
             let y_err_max = get Y 
                     (Array.mapi (fun i yi -> yi +. (err i yi)) h.H1.arr) y_min in
             (x,y,y_err_min,y_err_max)
         in
       (*  let all_dots_opt=get_dots (h::add_hists) in
         let all_dots_opt2=get_dots add_hists2 in
         let all_dots_opt3=get_dots add_hists3 in*)
         let plot ~opt h =
             def_style ~color:opt.color ();
(*             let pos = match opt.plot1d with 
                    Lines | Columns -> Left
                  | Polyline | Err_y -> Center
             in
             let x,y,y_err_min,y_err_max = 
                 get_dots pos h
             in
             ( match opt.plot1d with 
                 Lines -> draw_line (x,y) (*last_x*) x_max
               | Polyline -> draw_polyline (x,y)
               | Columns -> draw_bin (x,y)
               | Err_y -> draw_err (x,y,y_err_min,y_err_max)
             )*)
             ( match opt.plot1d with 
                 Lines | Columns ->
                     let x,y,y_err_min,y_err_max = get_dots Left h in
                     ( match opt.plot1d with 
                         Lines -> draw_line (x,y) x_max 
                       | Columns ->draw_bin (x,y))
               | Polyline | Err_y ->
                     let x,y,y_err_min,y_err_max = get_dots Center h in
                     ( match opt.plot1d with
                         | Polyline -> draw_polyline (x,y)
                         | Err_y -> draw_err (x,y,y_err_min,y_err_max))
               | Line_err -> 
                     let x,y,y_err_min,y_err_max = get_dots Left h in
                     draw_line (x,y) x_max;
                     let x,y,y_err_min,y_err_max = get_dots Center h in
                     draw_err (x,y,y_err_min,y_err_max)
             )

         in
(*         let rec list_iter2_opt ?(last_opt=draw_opt_1) f l lopt =
             ( match l with
               | [] -> ()
               | a::l -> 
                    ( match lopt with
                      | [] -> 
                              f a last_opt; list_iter2_opt ~last_opt f l []
                      | a2::lopt -> 
                              f a a2; list_iter2_opt ~last_opt:a2 f l lopt
                    )
             ) in*)
         let color=ref 1 in
         list_iter2_opt ~last_opt:draw_opt_1 (fun h opt ->
(*             def_style ~color:!color ();*)
             plot ~opt h; 
             incr color
         ) hists (draw_opt_1::draw_opt_other);
         list_iter2_opt ~last_opt:draw_opt_1 (fun h opt -> 
(*             def_style ~color:!color ();*)
(*             plwid 1;*)
             plot h ~opt; 
             incr color
         ) sec_hists sec_opt;
         draw_stats_for hists ~list_info;
    (*     if List.length add_hists3 > 0 then
             draw_stats_for ~place:Line2 add_hists3 ~list_info;*)
     );
         let wx,wy,n = plcalc_world 0.5 0.5 in 
         printf "!!!%f %f %i\n" wx wy n; flush stdout;
     plflush ()
end

module H2_pl (*: H2_draw*) = struct
 open Histo.H2

 let make_box plotopt sx sy =
     let opt s proj =
         if List.mem proj plotopt.grid
         then s^"gh"
         else s
     in
     plbox (opt sx X) 0. 0 (opt sy Y) 0. 0
 ;;

 let draw_box ~plot_opt h=
     let a=h.arr in
     let lx=Array.length a 
     and ly=Array.length a.(0)  in
     let (_,max)=Matrix.minfmaxf a in
     let s_x=(vp_x_max-.vp_x_min)/.(Axis.range h.xa)
     and s_y=(vp_y_max-.vp_y_min)/.(Axis.range h.ya) in
 (*    let gx_int=(vp_x_max-.vp_x_min)/.float lx 
     and gy_int=(vp_y_max-.vp_y_min)/.float ly in*)
     plwid 1;
     for ix=0 to lx-1 do
         for iy=0 to ly-1 do
           if Glob.is_finite a.(ix).(iy) && a.(ix).(iy) > 0.
           then (
                 let hw=0.47*.sqrt( a.(ix).(iy) /. max )  in
(*             if hw>0. then ( *)
                 (*let x=0.5 +. float ix and y=0.5 +. float iy in*)
		 let x=x_of_bin h ix and y=y_of_bin h iy in
                 let x_int=hw*.(interval h X ix) and y_int=hw*.(interval h Y iy) in
                (*let l=(x-.x_int) and r=(x+.x_int)
                 and d=(y-.y_int) and u=(y+.y_int) in*)
		 let l=x-.x_int and r=x+.x_int
                 and d=y-.y_int and u=y+.y_int in
                 plvpor vp_x_min vp_x_max vp_y_min vp_y_max;
                  plwind (lower h X) (upper h X) (lower h Y) (upper h Y);
		(* plwind (h.x_min-.hw) (h.x_max+.hw) (h.y_min-.hw) (h.y_max+.hw);*)
                 plcol1 0.8;
                 plfill [|l;l;r;r|] [|d;u;u;d|];
               (*let gxw=hw *. gx_int and gyw=hw *. gy_int in
                 let gx=vp_x_min +. (0.5 +. float ix) *. gx_int 
                 and gy=vp_y_min +. (0.5 +. float iy) *. gy_int in
                 plvpor (gx-.gxw) (gx+.gxw) (gy-.gyw) (gy+.gyw);
                 plwind 0. 1. 0. 1.;
                 plcol0 1;
                 plbox "bc" 0. 0 "bc" 0. 0;
	       *)
		 let gx old_coord = vp_x_min+.s_x*.(old_coord-.(lower h X))
		 and gy old_coord= vp_y_min+.s_y*.(old_coord-.(lower h Y)) in
		 let gl=gx l and gr=gx r and gd=gy d and gu=gy u in
		   plvpor gl gr gd gu;
		   plwind 0. 1. 0. 1.;
		   plcol0 1;
		   plbox "bc" 0. 0 "bc" 0. 0;
           )
         done;
     done;
     plvpor vp_x_min vp_x_max vp_y_min vp_y_max;
     plwind (lower h X) (upper h X) (lower h Y) (upper h Y);
(*     plbox "bcnst" 0. 0 "bcnstv" 0. 0;*)
     make_box plot_opt "bcnst" "bcnstv";
     pllab (Axis.name h.xa) (Axis.name h.ya) h.info.title;
     ()
 let level a max min =
     if max <> min then
          (a -. min) /. (max-.min)
      else 1.
     
 let draw_palette min max =
     plvpor (0.83*.vp_x_max) (0.95*.vp_x_max) vp_y_min vp_y_max;
     plwind 0. 1. min max;
     let n_colors=100 in
     let color_int=(max-.min)/.float n_colors in
     for l=0 to n_colors do
         let z=min+. color_int*.float l in
         let color_level = level z min max in
         let d=z in
         let u=d +. color_int in
         let l=0. and r=1. in
         plcol1 color_level;
         plfill [|l;l;r;r|] [|d;u;u;d|]
     done;
     plcol0 1;
     plbox "bc" 0. 0 "bcsmtv" 0. 0;
     ()

(* let get_coord (x,y) =*)
(*     let fr_x = (x -. vp_x_min) /. (0.81*.vp_x_max -. vp_x_min) in
     let fr_y = (y -. vp_y_min) /. (      vp_y_max -. vp_y_min) in*)


 let draw_shadebox ~palette h=
     plvpor vp_x_min (0.81*.vp_x_max) vp_y_min vp_y_max;
     plwind (lower h X) (upper h X) (lower h Y) (upper h Y);
     set_2d_palette palette;
     let a=h.arr in
     let lx=Array.length a 
     and ly=Array.length a.(0)  in
     let (min,max)=Matrix.minfmaxf a in
     for ix=0 to lx-1 do
         for iy=0 to ly-1 do
           if Glob.is_finite a.(ix).(iy) 
           then (
             let color_level=level a.(ix).(iy) min max in
             let l = x1_of_bin h ix in
             let r = x2_of_bin h ix in
             let d = y1_of_bin h iy in
             let u = y2_of_bin h iy in
             plcol1 color_level;
             plfill [|l;l;r;r|] [|d;u;u;d|]
           )
         done;
     done;
     plcol0 1;
     plbox "bcnst" 0. 0 "bcnstv" 0. 0;
     pllab (Axis.name h.xa) (Axis.name h.ya) h.info.title;
     draw_palette min max;
     ()
 let draw_contour ~plot_opt h=
     next_page (lower h X) (upper h X) (lower h Y) (upper h Y);
     make_box plot_opt "bcnst" "bcnstv";
(*     plbox "bcnst" 0.0 0 "bcnstv" 0.0 0;*)
     pllab "" "" h.info.title;
     let min,max=Matrix.minfmaxf h.arr in
     let n_colors=7 in
     let color_interval=(max-.min)/.float n_colors in
     let z_colors=Array.init n_colors (
         fun i -> min+.color_interval*.float i) in
     plset_pltr (fun ix iy -> xw_of_bin h ix, yw_of_bin h iy);
     let nx=Array.length h.arr 
     and ny=Array.length h.arr.(0) in
     plcont h.arr 1 (nx-1) 1 (ny-1) z_colors;
(*     plshades h.arr h.x_min h.x_max h.y_min h.y_max z_colors 0 1 1 true;*)
(*     plshades h.arr h.x_min h.x_max h.y_min h.y_max z_colors 0 1 1 1;*)
     ()
 let list_info h=
     let (_,(mean_x,rms_x),(mean_y,rms_y))=max_mean_rms h in
     [
         ("N=", Int64.to_float h.entries);
        (* ("∫(range)=", integral h);*)
         ("∑#d∈☐#uz=", integral h);
         ("∑#d∉☐#uz=", Matrix.sum h.out);
         ("Mx=", mean_x);
         ("My=", mean_y);
         ("√Dx=", rms_x);
         ("√Dy=", rms_y);
     ] 

 let get_log_z h =
     let new_arr = Matrix.map log h.arr
     in {h with arr=new_arr}

 let draw_hist ~plot_opt ~show_opt hl =
     let opt1 = match show_opt with [] -> standard_show | opt::_ -> opt in
     let plot h' opt = (
         let h = if List.mem Log_Z plot_opt.logscale then get_log_z h' else h' in
         def_style ~color:opt.color ();
         (match opt1.plot2d with 
             Box -> draw_box ~plot_opt h
           | Grayscale -> draw_shadebox ~palette:1 h
           | Color -> draw_shadebox ~palette:2 h
           | Contour -> draw_contour ~plot_opt h
         ) 
     ) in
     list_iter2_opt ~last_opt:opt1 plot hl show_opt;
(*	 List.iter plot hl;*)
     draw_stats_for hl ~list_info;
     plflush ()
end

let finalize () =
    plend ()

let draw_hist 
        ?(plot_opt = standard_plot) ?(opt=[]) ?(sec_hists=[]) ?(sec_opt=[]) 
        hists =
    (match hists with 
    | [] -> ()
    | (H1d h)::_ -> 
          let filter l = List.map get_1d l in
          H1_pl.draw (filter hists) ~plot_opt
                ~opt ~sec_hists:(filter sec_hists) ~sec_opt
    | (H2d h)::_ -> 
          let filter l = List.map get_2d l in
          H2_pl.draw_hist (filter hists) ~plot_opt ~show_opt:opt
    )

let init (nx, ny) =
    pl_setcontlabelformat 4 3;
    pl_setcontlabelparam 0.015 0.5 0.1 1;
    cmap1_init ();
    plspal0 "cmap0_black2.pal";
    plssub nx ny;
    ()

let advance (x,y) (nx,ny) =
    if x >= nx || y >= ny || x<0 || y<0
    then failwith "plplot incorrect sub-range"
    else (
(*        Printf.printf "advance to %i,%i => %i sub-page\n" x y (1+x+y*nx);*)
	flush stdout;
        pladv (1+x+y*nx); 
   )

let init_for_printing ?(color=false) p =
    init p;
(*    plspal1 "cmap1_gray.pal" true;*)
    if color then (
        plspal0 "cmap0_black2.pal"
    ) else ( 
        plspal0 "cmap0_black_on_white.pal"
    )
;;

let end_draw () =
    Plplot.plend ()

(** save plot (copy-paste from plplot ocaml examples) to pdf-file *)
let save_plot_copy ~device ~name =
  let cur_strm = plgstrm () in (* get current stream *)
  let _ = plmkstrm () in  (* create a new one *)

  plsdev device; (* new device type. Use a known existing driver *)
  plsfnam name; (* file name *)

  plcpstrm cur_strm false; (* copy old stream parameters to new stream *)
  plreplot (); (* do the save *)
  plend1 (); (* close new device *)

  plsstrm cur_strm; (* and return to previous one *)
  ()

exception Unknown_file_type of string

let save_plot name = 
  let file_has suffix = Filename.check_suffix name suffix in
  let formats = [ ".pdf","pdfcairo"; ".ps","pscairo"; ".eps","pstex"; 
        ".png","pngcairo"; ] in
  let (_,device) = try (
      List.find (fun (suffix,_) -> file_has suffix) formats
  ) with Not_found -> raise (Unknown_file_type name)
  in
  save_plot_copy ~device ~name

let test_hist () =
    let hi=init_hist (H1.h ~xa:{n=100;x1= -5.;x2=5.}
    {name="h";title="gaus"}) (Glob.gaus
    1.0 0.5 0.5) in
    init (1,1);
    H1_pl.draw [hi]

module Distinct=struct
    let init () =
        (*plparseopts [|"open";"-geometry";"600x500"|] [PL_PARSE_FULL];*)
        plparseopts [|"open";"-geometry";"1600x1100"|] [PL_PARSE_FULL];
        Plplot.plsdev "xcairo";
        Plplot.plinit ();
        init (1,1);
        ()
    let finish () =
        plflush ();
        plend ()
    let draw_hist hists ?opt ?sec_hists ?sec_opt =
        init ();
        draw_hist hists ?opt ?sec_hists ?sec_opt;
        finish ();
end

